#!/usr/bin/python3
# -*- coding: utf-8 -*-

import os
import sys
import cv2
import json
import tqdm
import glob
from PIL import Image
import xml.etree.ElementTree as ET


def load_voc_annotations(categories, anno_folder, file_id_map):
    """ 加载voc格式的标注数据
    """

    assert(len(file_id_map) > 0)

    annotations = {}

    for xml_name in tqdm.tqdm(file_id_map):
        xml_name = xml_name.strip()

        annotations[xml_name] = {}
        annotations[xml_name]["gt"] = []
        annotations[xml_name]["dt"] = []

        xml_abs_name = os.path.join(anno_folder, xml_name+".xml")
        if not os.path.exists(xml_abs_name):
            continue
        tree = ET.parse(xml_abs_name)
        root = tree.getroot()
        
        # size = root.find('size')
        width = 10000#int(size.find('width').text)
        height = 10000#int(size.find('height').text)
        if width <=0 or height <= 0:
            print("find width or height less than 0 in {}".format(xml_name))
            continue

        for obj in root.iter('object'):
            cls_name = obj.find('name').text
            if cls_name not in categories:
                print("ANNOTATIONS LOAD: can not recognize class name {}".format(cls_name))
                continue

            xmlbox = obj.find('bndbox')
            bbox = [int(xmlbox.find('xmin').text),
                    int(xmlbox.find('ymin').text),
                    int(xmlbox.find('xmax').text),
                    int(xmlbox.find('ymax').text)]
            bbox[0] = max(0, bbox[0])
            bbox[1] = max(0, bbox[1])
            bbox[2] = min(width-1, bbox[2])
            bbox[3] = min(height-1, bbox[3])
            assert(bbox[2]>bbox[0] and bbox[3]>bbox[1])

            annotations[xml_name]["gt"].append({"bbox": bbox, "category": cls_name})

    return annotations


def load_annotations(categories, anno_folder, file_id_map):
    """加载标注数据，voc标注格式
    输入：文件或者文件夹列表
    输出：
    [
        {
            "image_name": {
                "gt": [anno1, anno2, ...],
                "dt": []
            },
            ...
        }
    ]

    anno1格式：
    {
        "bbox": [xmin, ymin, xmax, ymax],
        "category": cls_name,
    }
    """

    print("loading annotations ")

    annotations = load_voc_annotations(categories, anno_folder, file_id_map)

    for image_name in annotations:
        for obj in annotations[image_name]["gt"]:
            categories[obj["category"]]["gtn"] += 1


    gt_images_cnt = 0
    gt_objects_cnt = 0
    for name in annotations:
        if len(annotations[name]["gt"]) > 0:
            gt_images_cnt += 1
            gt_objects_cnt += len(annotations[name]["gt"])
    print("load {} annotations from {} images".format(gt_objects_cnt, gt_images_cnt))

    return categories, annotations

